#ifndef CKKSNTT_NUMB_H_
#define CKKSNTT_NUMB_H_

#include <set>

#include <arm_neon.h>
#include <ctime>
#include <iostream>

#include <stdio.h>
#include <stdint.h>
#include <stdexcept>
#include <vector>
#include <sys/time.h>
#include <string>
#include <math.h>

using namespace std;

/*
 * @brief 计算给定整数的负值。
 *
 * 该函数接受一个 64 位无符号整数 `a`，计算其负值，并将结果赋给引用参数 `r`。
 * 此操作相当于对 `a` 取补并加一。
 *
 * @param r 结果变量,输出无符号整数，保存 `a` 的负值。
 * @param a 输入无符号整数，表示需要求负值的数。
 */
inline void Negate(uint64_t &r, uint64_t a)
{
    r = -a;
}

/*
 * @brief 对无符号整数进行条件减法操作。
 *
 * 该函数首先计算 `a - q`，然后根据结果的符号来决定是否将 `q` 加回到 `a`。
 * 如果 `a - q` 的结果为负（即 `a` 小于 `q`），则通过加回 `q` 来修正 `a` 的值。
 *
 * @param a 输入并更新的无符号整数。该值在函数执行后会被修改为根据条件计算后的结果。
 * @param q 要减去的无符号整数 `q`。
 */
inline void ComSubQ(uint64_t &a, uint64_t q)
{
    uint64_t tmp = a - q;
    a = tmp + (tmp >> 63) * q;
}
/*
 * @brief 对 NEON 向量进行条件减法操作。
 *
 * 该函数使用 NEON SIMD 指令对 64 位无符号整数向量进行条件减法。如果 `a[i] >= q[i]`，则执行 `a[i] - q[i]`，否则不改变
 * `a[i]`。
 *
 * @param a 输入并更新的 64 位无符号整数向量。该向量在函数执行后会被修改为根据条件计算后的结果。
 * @param q 要减去的 64 位无符号整数向量 `q`。
 *
 * @note 该操作会修改 `a` 向量的值，并且适用于 ARM 架构的 SIMD 指令集。
 */
inline void ComSubQNeon(uint64x2_t &a, uint64x2_t q)
{
    uint64x2_t mask = vcgeq_u64(a, q);         // mask = (a >= q) ? 0xFFFFFFFFFFFFFFFF : 0x0
    uint64x2_t sub_value = vandq_u64(mask, q); // mask AND q (preserve q where condition met)
    a = vsubq_u64(a, sub_value);               // a = a - sub_value (only subtract where condition met)
}

/*
 * @brief 对两个无符号整数进行模加操作。
 *
 * 该函数先对输入的 `a` 和 `b` 分别执行条件减法（确保它们小于模数
 * `p`），然后对它们进行加法操作，最后检查和是否大于或等于模数 `p`。 如果和大于等于模数 `p`，则将其减去模数
 * `p`，实现模运算。
 *
 * @param r 结果值，保存加法后的模结果。该值会被更新为 `(a + b) % p`。
 * @param a 第一个加数，无符号整数。
 * @param b 第二个加数，无符号整数。
 * @param p 模数，无符号整数。
 */
inline void AddMod(uint64_t &r, uint64_t a, uint64_t b, uint64_t p)
{
    ComSubQ(a, p);
    ComSubQ(b, p);
    r = a + b;
    if (r >= p) {
        r -= p;
    }
}

/*
 * @brief 使用汇编Barrett算法进行128位取模运算。
 *
 * 该函数使用Barrett算法计算一个128位数 `a` 对一个64位模数 `q` 进行模运算，具体过程通过使用乘法和取模优化来提高效率。
 * 函数使用内联汇编实现，利用多个乘法和加法步骤计算模运算。
 *
 * @param a 128位数 `a`，需要进行模运算的数。函数将修改此参数，最终返回模运算结果的低64位。
 * @param q 64位模数 `q`。
 * @param mu_high Barrett算法中的高64位乘法常数。
 * @param mu_low Barrett算法中的低64位乘法常数。
 *
 * @return 返回模运算结果的低64位。即计算 `(a * mu) % q` 的结果。
 */
inline uint64_t Barrett128(__uint128_t &a, uint64_t q, uint64_t mu_high, uint64_t mu_low)
{
    uint64_t a_high = a >> 64;
    uint64_t a_low = static_cast<uint64_t>(a);
    uint64_t qq = 2 * q;
    asm volatile(
        // Step 1: Calculate a_low * mu_high, store high 64 bits in x1
        "umulh x1, %[a_low], %[mu_high]\n\t"

        // Step 2: Calculate a_high * mu_low and add to x1
        "umulh x0, %[a_high], %[mu_low]\n\t"
        "add x1, x1, x0\n\t"

        // Step 3: Calculate a_high * mu_high and add to x1
        "mul x0, %[a_high], %[mu_high]\n\t"
        "add x1, x1, x0\n\t"

        // Step 4: Calculate x1 * q
        "mul x1, x1, %[q]\n\t"

        // Step 5: Calculate a_low - x1
        "sub x1, %[a_low], x1\n\t"

        // Step 6: if (x1 >= qq) x1 -= qq
        "subs x0, x1, %[qq]\n\t"  // Calculate x1 - qq and set flags
        "csel x1, x0, x1, hs\n\t" // If x1 >= qq (hs), select x1 - qq, else select x1

        // Step 7: if (x1 >= q) x1 -= q
        "subs x0, x1, %[q]\n\t"         // Calculate x1 - q and set flags
        "csel %[a_low], x0, x1, hs\n\t" // If x1 >= q (hs), select x1 - q, else select x1

        // Store the result back to a_low
        // "mov %[a_low], x1\n\t"

        : [ a_low ] "+r"(a_low) // Output: a_low (modified)
        : [ a_high ] "r"(a_high), [ mu_low ] "r"(mu_low), [ mu_high ] "r"(mu_high), [ q ] "r"(q), [ qq ] "r"(qq)
        : "x0", "x1", "memory", "cc");

    return a_low;
}

/*
 * @brief 单次使用汇编Barrett算法进行128位取模运算的封装函数。
 *
 * @param a 128位数 `a`，需要进行模运算的数。函数将修改此参数，最终返回模运算结果的低64位。
 * @param q 64位模数 `q`。
 * @param muHigh Barrett算法中的高64位乘法常数。
 * @param muLow Barrett算法中的低64位乘法常数。
 *
 * @return 返回模运算结果的低64位，即计算 `a % q` 的结果。
 */
inline uint64_t SingleBarrettNew(__uint128_t &a, uint64_t q, uint64_t muHigh, uint64_t muLow)
{
    return Barrett128(a, q, muHigh, muLow);
}

/*
 * @brief Barrett算法进行64位数的模运算。
 *
 * 该函数使用Barrett算法执行64位数 `a` 对64位模数 `q` 的模运算。Barrett算法的主要目的是通过预先计算常数 `mu`，
 * 来加速大数模运算。该函数通过以下步骤实现：
 * 1. 计算 `a * mu`。
 * 2. 计算 `a - (a * mu) * q`。
 * 3. 如果结果大于等于 `q`，则减去 `q`，以确保结果在模数范围内。
 *
 * @param[in,out] a 64位数 `a`，需要进行模运算的数。函数将修改此参数，最终返回模运算结果。
 * @param[in] q 64位模数 `q`，用于模运算。
 * @param[in] mu Barrett算法中的乘法常数，用于加速计算。
 *
 * @return 返回 `a % q` 的结果。
 *
 * @note 本函数使用内联汇编实现Barrett算法，并进行高效的模运算。它直接修改输入参数 `a`，并返回计算结果。
 */
inline uint64_t Barrett64(uint64_t a, uint64_t q, uint64_t mu)
{
    asm volatile(
        // Step 1: Calculate a * mu, store high 64 bits in x1
        "umulh x1, %[a], %[mu]\n\t"

        // Step 2: Calculate x1 * q
        "mul x1, x1, %[q]\n\t"

        // Step 3: Calculate a - x1
        "sub x1, %[a], x1\n\t"

        // Step 4: if (x1 >= q) x1 -= q
        "subs x0, x1, %[q]\n\t"     // Calculate x1 - q and set flags
        "csel %[a], x0, x1, hs\n\t" // If x1 >= q (hs), select x1 - q, else select x1

        : [ a ] "+r"(a)                // Output: a (modified)
        : [ q ] "r"(q), [ mu ] "r"(mu) // Inputs: q, mu
        : "x0", "x1", "cc"             // Clobbers: x0, x1, condition flags
    );
    return a;
}

/*
 * @brief 使用NEON指令集进行Barrett算法的模运算（64位）。
 * 函数通过内联汇编实现，分别对 `a` 和 `q` 的低32位进行运算。
 *
 * @param[out] res 结果向量，保存计算后的模运算结果。
 * @param[in] a 64位数 `a`，需要进行模运算的数。此为一个64位的向量。
 * @param[in] q 64位模数 `q`，用于模运算。
 * @param[in] mu Barrett算法中的常数 `mu`，用于加速模运算。
 *
 * @return 无返回值，结果直接通过 `res` 输出。
 */
inline void Barrett64_NEON(uint64x2_t *res, uint64x2_t a, uint64x2_t q, uint64x2_t mu)
{
    uint64_t res0 = 0, res1 = 0;

    asm volatile("umulh x1, %[a0], %[mu0]\n\t"
        "mul x1, x1, %[q0]\n\t"
        "sub x1, %[a0], x1\n\t"
        "subs x0, x1, %[q0]\n\t"
        "csel %[res0], x0, x1, hs\n\t"

        : [ res0 ] "=r"(res0)
        : [ a0 ] "r"(vgetq_lane_u64(a, 0)), [ q0 ] "r"(vgetq_lane_u64(q, 0)), [ mu0 ] "r"(vgetq_lane_u64(mu, 0))
        : "x0", "x1", "cc");

    asm volatile("umulh x1, %[a1], %[mu1]\n\t"
        "mul x1, x1, %[q1]\n\t"
        "sub x1, %[a1], x1\n\t"
        "subs x0, x1, %[q1]\n\t"
        "csel %[res1], x0, x1, hs\n\t"

        : [ res1 ] "=r"(res1)
        : [ a1 ] "r"(vgetq_lane_u64(a, 1)), [ q1 ] "r"(vgetq_lane_u64(q, 1)), [ mu1 ] "r"(vgetq_lane_u64(mu, 1))
        : "x0", "x1", "cc");

    *res = vsetq_lane_u64(res0, *res, 0);
    *res = vsetq_lane_u64(res1, *res, 1);
}

/*
 * @brief 对两个数进行模减法运算。
 *
 * @param r 结果，保存计算后的模减法结果。
 * @param a 被减数 `a`。
 * @param b 减数 `b`。
 * @param m 模数 `m`，用于模减法运算。
 *
 */
inline void SubMod(uint64_t &r, uint64_t a, uint64_t b, uint64_t m)
{
    if (m == 0) {
        cout << "模数为0!" << endl;
    }
    ComSubQ(a, m);
    ComSubQ(b, m);
    r = b % m;
    r = (a + m - r) % m;
}

/*
 * @brief 对两个数进行模乘法运算。
 *
 * @param r 结果，保存计算后的模乘法结果。
 * @param a 被乘数 `a`。
 * @param b 乘数 `b`。
 * @param m 模数 `m`，用于模乘法运算。
 */
inline void MulMod(uint64_t &r, uint64_t a, uint64_t b, uint64_t m)
{
    if (m == 0) {
        cout << "模数为0!" << endl;
    }
    if (a > m)
        a = a % m;
    if (b > m)
        b = b % m;
    unsigned __int128 mul = static_cast<unsigned __int128>(a) * (b);
    mul %= static_cast<unsigned __int128>(m);
    r = static_cast<uint64_t>(mul);
}

/*
 * @brief 使用汇编 Barrett 约减算法计算模乘法：a * b mod p
 *
 * 该函数通过 Barrett 约减算法高效地计算给定的两个数 `a` 和 `b` 在模数 `p` 下的乘积 `(a * b) % p`。
 * 首先通过 ASM 汇编计算 `a * b` 的乘积，并将结果经过 Barrett 约减处理，最终得到模 `p` 下的结果。
 *
 * @param[out] r 计算结果，存储 `(a * b) % p`。
 * @param[in] a 被乘数 `a`。
 * @param[in] b 乘数 `b`。
 * @param[in] p 模数 `p`。
 * @param[in] twop 模数 `p` 的两倍 `2 * p`，用于减少溢出。
 * @param[in] prHigh Barrett 预计算常量的高 64 位。
 * @param[in] prLow Barrett 预计算常量的低 64 位。
 */
inline void MulModBarrett(uint64_t &r, uint64_t a, uint64_t b, uint64_t p, uint64_t twop, uint64_t prHigh,
    uint64_t prLow)
{
    if (p == 0) {
        cout << "模数为0!" << endl;
    }
    if (a > p)
        a = a % p;
    if (b > p)
        b = b % p;
    asm volatile(
        // 计算 a*b  低位存x3 高位存x3寄存器
        "mul   x3 ,%[a] , %[b] \n\t "
        "umulh x4 ,%[a] , %[b] \n\t"
        // 计算Mul_l 与预计算量pr的乘积，高64位存在res中
        "umulh x1, x3, %[mu_high]\n\t"

        // 计算 a.high * mu.low + res : 先计算高64位乘法，再累加res,要存放到res中
        "umulh x0, x4, %[mu_low]\n\t"
        "add x1, x1, x0\n\t"

        // 计算 a.high * mu.high + res,先计算低64位乘法， 再累加res,要存放到res中
        "mul x0, x4, %[mu_high]\n\t"
        "add x1, x1, x0\n\t"

        // 计算 res * q 存回res
        "mul x1, x1, %[q]\n\t"

        // 计算 a.low - res,存回res
        "sub x1, x3, x1\n\t"

        // 计算res 和2*q的比较，如果直接减发生溢出就set flag
        "subs x0, x1, %[qq]\n\t"
        "csel x1, x0, x1, hs\n\t"

        // 计算res 和q的比较，如果直接减发生溢出就set flag
        "subs x0, x1, %[q]\n\t"
        "csel %[r], x0, x1, hs\n\t"
        // Store the result back to a_low
        //   "mov %[r], x1\n\t"
        : [ a ] "+r"(a), [ b ] "+r"(b), [ r ] "+r"(r)
        : [ mu_low ] "r"(prLow), [ mu_high ] "r"(prHigh), [ q ] "r"(p), [ qq ] "r"(twop)
        : "x0", "x1", "x3", "x4", "memory", "cc");
}
/*
 * @brief 使用 Barrett 约减算法计算 a mod p
 *
 * 该函数使用 Barrett 约减算法计算给定的数字 `a` 在模数 `p` 下的结果 `(a % p)`。
 *
 * @param[out] r 计算结果，存储 `a % p`。
 * @param[in] a 被取模数 `a`，可以是 `uint64_t` 或 `unsigned __int128` 类型。
 * @param[in] p 模数 `p`。
 * @param[in] twop 模数 `p` 的两倍 `2 * p`，用于减少溢出。
 * @param[in] prHigh Barrett 预计算常量的高 64 位。
 * @param[in] prLow Barrett 预计算常量的低 64 位。
 */
inline void ModBarrett(uint64_t &r, uint64_t a, uint64_t p, uint64_t twop, uint64_t prHigh, uint64_t prLow)
{
    __uint128_t a_ = static_cast<__uint128_t>(a);
    uint64_t a_high = static_cast<uint64_t>(a_ >> 64);
    uint64_t a_low = static_cast<uint64_t>(a_);
    r = Barrett128(a_, p, prHigh, prLow);
}
inline void ModBarrett(uint64_t &r, unsigned __int128 a, uint64_t p, uint64_t twop, uint64_t prHigh, uint64_t prLow)
{
    uint64_t a_high = a >> 64;
    uint64_t a_low = static_cast<uint64_t>(a);
    r = Barrett128(a, p, prHigh, prLow);
}
/*
 * @brief 使用汇编 Shoup 模运算优化算法计算 a * b mod q
 *
 * 该函数使用 Shoup 模运算优化算法计算给定的数字 `a` 和 `b` 的乘积，然后在模 `q` 下取模。
 *
 * @param[out] res 计算结果，存储 `(a * b) % q`。
 * @param[in] a 操作数 `a`。
 * @param[in] b 操作数 `b`。
 * @param[in] bShoup 预计算的shoup参数。
 * @param[in] q 模数 `q`。
 */
inline void MulModShoup(uint64_t &res, uint64_t a, uint64_t b, uint64_t bShoup, uint64_t q)
{
    // 处理第一个元素
    asm volatile("mul    x3, %[a], %[bShoup] \n\t" // x3 = a[0] * bShoup[0] (low 64 bits)
        "umulh  x4, %[a], %[bShoup] \n\t"          // x4 = a[0] * bShoup[0] (high 64 bits)
        "mul    x1, x4, %[q]        \n\t"          // x1 = (a[0] * bShoup[0] high) * q[0]
        "mul    x0, %[a], %[b]     \n\t"           // x0 = a[0] * b[0] (low 64 bits)
        "sub    x1, x0, x1           \n\t"         // x1 = (a[0] * b[0]) - ((a[0] * bShoup[0] high) * q[0])
        "subs   x0, x1, %[q]        \n\t"          // x0 = x1 - q[0]
        "csel   %[res], x1, x0, lo  \n\t"          // res0 = (x1 < q[0]) ? x1 : x0
        : [ res ] "=r"(res)
        : [ a ] "r"(a), [ b ] "r"(b), [ bShoup ] "r"(bShoup), [ q ] "r"(q)
        : "x0", "x1", "x3", "x4", "cc");
}
/*
 * @brief 计算两个数的最大公约数（GCD）
 *
 * 该函数使用欧几里得算法计算给定两个非负整数 `a` 和 `b` 的最大公约数。欧几里得算法基于以下性质：
 * - gcd(a, b) = gcd(b % a, a)
 * 即：两个数的最大公约数等于其中一个数除以另一个数的余数与另一个数的最大公约数。
 *
 * @param[in] a 操作数 `a`。
 * @param[in] b 操作数 `b`。
 *
 * @return 返回 `a` 和 `b` 的最大公约数。
 */
inline uint64_t Gcd(uint64_t a, uint64_t b)
{
    if (a == 0) {
        return b;
    }
    return Gcd(b % a, a);
}

inline long Gcd(long a, long b)
{
    if (a == 0) {
        return b;
    }
    return Gcd(b % a, a);
}

/*
 * @brief ：计算 x^y 对模 modulus 取模的结果。
 *
 * @param x：底数 x。
 * @param y：指数 y。
 * @param modulus：模数 modulus。
 * @return uint64_t：返回 x^y 对模 modulus 取模的结果。
 */
inline uint64_t PowMod(uint64_t x, uint64_t y, uint64_t modulus)
{
    uint64_t res = 1;
    while (y > 0) {
        if ((y & 1) == 1) {
            MulMod(res, res, x, modulus);
        }
        y = y >> 1;
        MulMod(x, x, x, modulus);
    }
    return res;
}
/*
 * @brief 计算整数 a 在模 m 下的逆元
 *
 * 使用扩展欧几里得算法计算模逆元，即找到一个整数 x 使得：
 *
 * @param a 要计算逆元的整数
 * @param m 模数
 * @return 如果逆元存在，返回逆元；否则返回 -1 表示不存在逆元。
 */
inline uint64_t InvMod(uint64_t a, uint64_t m)
{
    int64_t m0 = m;
    int64_t x0 = 0, x1 = 1;
    int64_t q, t;

    if (m == 1)
        return 0; // 任何数 mod 1 都没有逆元

    while (a > 1) {
        // q 是商
        q = a / m;
        t = m;
        // m 是余数，a 是被除数
        m = a % m;
        a = t;
        // 更新 x0 和 x1
        t = x0;
        x0 = x1 - q * x0;
        x1 = t;
    }

    // 确保 x1 是正数
    if (x1 < 0)
        x1 += m0;

    if (x1 == -1) {
        //        throw invalid_argument("Inverse doesn't exist!!!");
    }

    return x1;
}

/*
 * @brief ：计算 x 在模 m 下的逆元素（适用于 uint32_t 类型）。
 *
 * @param x：输入值 x。
 * @param m：模数 m。
 * @return uint64_t：返回 x 在模 m 下的逆元素。
 */
inline uint64_t InvMod(uint32_t x, uint32_t m)
{
    if (m == 0) {
        cout << "模数为0 " << endl;
    }
    uint64_t temp = x % m;
    if (Gcd(temp, m) != 1) {
        //        throw invalid_argument("Inverse doesn't exist!!!");
    } else {
        return PowMod(temp, m - 2, m);
    }
}

/*
 * @brief 计算整数 x 的 y 次幂
 *
 * 该函数使用快速幂算法（Exponentiation by Squaring）计算 x 的 y 次幂。
 * 每次将底数平方，指数右移一位来减少问题规模，直到指数为 0 为止。
 *
 * @param x 底数
 * @param y 指数
 * @return 返回 x 的 y 次幂，类型为 uint64_t
 */
inline uint64_t Pow(uint64_t x, uint64_t y)
{
    uint64_t res = 1;
    while (y > 0) {
        if (y & 1) {
            res *= x;
        }
        y = y >> 1;
        x *= x;
    }
    return res;
}
/*
 * @brief ：计算 x 的逆元素（通过快速幂算法）。
 *
 * @param x：输入值 x。
 * @return uint64_t：返回 x 的逆元素。
 */
inline uint64_t Inv(uint64_t x)
{
    return Pow(x, static_cast<uint64_t>(-1));
}

/*
 * @brief 反转一个 32 位整数的二进制位
 *
 * 该函数通过一系列位操作，将给定的 32 位整数的二进制位进行反转。例如，
 * 输入的二进制数 00000000000000000000000000000001 会变成 10000000000000000000000000000000。
 *
 * @param x 要反转的 32 位整数
 * @return 返回反转后的 32 位整数
 */
inline uint32_t BitReverse(uint32_t x)
{
    x = (((x & 0xaaaaaaaa) >> 1) | ((x & 0x55555555) << 1));
    x = (((x & 0xcccccccc) >> 2) | ((x & 0x33333333) << 2));
    x = (((x & 0xf0f0f0f0) >> 4) | ((x & 0x0f0f0f0f) << 4));
    x = (((x & 0xff00ff00) >> 8) | ((x & 0x00ff00ff) << 8));
    return ((x >> 16) | (x << 16));
}

/*
 * @brief ：将给定的整数分解为其所有的素因数
 *
 * @param s：存储素因数的集合
 * @param number：待分解的整数
 */
inline void FindPrimeFactors(set<uint64_t> &s, uint64_t number)
{
    while (number % 2 == 0) {
        s.insert(2);
        number /= 2;
    }
    for (uint64_t i = 3; i < sqrt(number); i++) {
        while (number % i == 0) {
            s.insert(i);
            number /= i;
        }
    }
    if (number > 2) {
        s.insert(number);
    }
}

/*
 * @brief ：查找给定模数的原根
 *
 * @param modulus：模数
 * @return uint64_t：模数的原根，如果不存在则返回 -1
 */
inline uint64_t FindPrimitiveRoot(uint64_t modulus)
{
    set<uint64_t> s;
    uint64_t phi = modulus - 1;
    FindPrimeFactors(s, phi);
    for (uint64_t r = 2; r <= phi; r++) {
        bool flag = false;
        for (auto it = s.begin(); it != s.end(); it++) {
            if (PowMod(r, phi / (*it), modulus) == 1) {
                flag = true;
                break;
            }
        }
        if (flag == false) {
            return r;
        }
    }
    return -1;
}

/*
 * @brief ：查找模数下的 M 次单位根
 *
 * @param M：单位根的阶
 * @param mod：模数
 * @return uint64_t： M 次单位根的值，如果无法找到则返回 -1
 */
inline uint64_t FindMthRootOfUnity(uint64_t M, uint64_t mod)
{
    uint64_t res;
    res = FindPrimitiveRoot(mod);
    if (M == 0) {
        cout << "模数为0 " << endl;
    }
    if ((mod - 1) % M == 0) {
        uint64_t factor = (mod - 1) / M;
        res = PowMod(res, factor, mod);
        return res;
    } else {
        return -1;
    }
}

/*
 * @brief ：测试一个数是否为素数
 *
 * @param p：待测试的数
 * @return bool：如果是素数返回 true，否则返回 false
 */
inline bool PrimeTest(uint64_t p)
{
    if (p < 2) {
        return false;
    }
    if (p != 2 && p % 2 == 0) {
        return false;
    }
    uint64_t s = p - 1;
    while (s % 2 == 0) {
        s /= 2;
    }
    for (long i = 0; i < 200; i++) {
        uint64_t temp1 = rand();
        temp1 = (temp1 << 32) | rand();
        temp1 = temp1 % (p - 1) + 1;
        uint64_t temp2 = s;
        uint64_t mod = PowMod(temp1, temp2, p);
        while (temp2 != p - 1 && mod != 1 && mod != p - 1) {
            MulMod(mod, mod, mod, p);
            temp2 *= 2;
        }
        if (mod != p - 1 && temp2 % 2 == 0) {
            return false;
        }
    }
    return true;
}

/*
 * @brief ：计算给定整数 x 对 p 的 Shoup 表示
 *
 * @param x：整数
 * @param p：模数
 * @return uint64_t： x 对 p 的 Shoup 表示
 */
inline uint64_t x_Shoup(uint64_t x, uint64_t p)
{
    if (p == 0) {
        cout << "模数为0 " << endl;
    }
    unsigned __int128 temp = x;
    temp <<= 64;
    return uint64_t(temp / p);
}
/*
 * @brief ：使用Barrett乘法算法对输入的多项式系数进行乘法并取模。
 *
 * 该函数通过NEON指令集实现Barrett乘法优化，执行大整数乘法并模一个指定的数
 * p。计算过程中使用了高效的硬件指令集，以加速模乘操作。
 *
 * @param r： 输出结果向量，存储计算后的模乘结果。
 * @param a： 输入的第一个多项式系数，表示待乘的第一个数。
 * @param b： 输入的第二个多项式系数，表示待乘的第二个数。
 * @param p： 用于取模的模数 p。
 * @param twop： 模数 p 的两倍值，用于进一步优化模操作。
 * @param prHigh： 预计算的模数 p 的高位部分，用于优化乘法运算。
 * @param prLow： 预计算的模数 p 的低位部分，用于优化乘法运算。
 */
inline void MulModBarrett_NEON(uint64x2_t *r, uint64x2_t a, uint64x2_t b, uint64x2_t p, uint64x2_t twop,
    uint64x2_t prHigh, uint64x2_t prLow)
{
    uint64_t res0 = 0, res1 = 0;

    // 处理第一个元素
    asm volatile(
        // 计算 a[0] * b[0] 低位存 x3, 高位存 x4
        "mul x3, %[a0], %[b0]\n\t"
        "umulh x4, %[a0], %[b0]\n\t"

        // 计算 Mul_l 与预计算量 prHigh 的乘积，高64位存 x1
        "umulh x1, x3, %[prHigh0]\n\t"

        // 计算 a[0].high * prLow + res (高位乘法再累加到 x1)
        "umulh x0, x4, %[prLow0]\n\t"
        "add x1, x1, x0\n\t"

        // 计算 a[0].high * prHigh + res (低位乘法再累加到 x1)
        "mul x0, x4, %[prHigh0]\n\t"
        "add x1, x1, x0\n\t"

        // 计算 x1 * p 存回 x1
        "mul x1, x1, %[p0]\n\t"

        // 计算 a[0].low - x1, 存回 x1
        "sub x1, x3, x1\n\t"

        // 计算 x1 和 2 * p 的比较，如果直接减发生溢出就 set flag
        "subs x0, x1, %[twop0]\n\t"
        "csel x1, x0, x1, hs\n\t"

        // 计算 x1 和 p 的比较，如果直接减发生溢出就 set flag
        "subs x0, x1, %[p0]\n\t"
        "csel %[res0], x0, x1, hs\n\t"

        : [ res0 ] "=r"(res0)
        : [ a0 ] "r"(vgetq_lane_u64(a, 0)), [ b0 ] "r"(vgetq_lane_u64(b, 0)), [ p0 ] "r"(vgetq_lane_u64(p, 0)),
        [ twop0 ] "r"(vgetq_lane_u64(twop, 0)), [ prHigh0 ] "r"(vgetq_lane_u64(prHigh, 0)),
        [ prLow0 ] "r"(vgetq_lane_u64(prLow, 0))
        : "x0", "x1", "x3", "x4", "cc");

    // 处理第二个元素
    asm volatile(
        // 计算 a[1] * b[1] 低位存 x3, 高位存 x4
        "mul x3, %[a1], %[b1]\n\t"
        "umulh x4, %[a1], %[b1]\n\t"

        // 计算 Mul_l 与预计算量 prHigh 的乘积，高64位存 x1
        "umulh x1, x3, %[prHigh1]\n\t"

        // 计算 a[1].high * prLow + res (高位乘法再累加到 x1)
        "umulh x0, x4, %[prLow1]\n\t"
        "add x1, x1, x0\n\t"

        // 计算 a[1].high * prHigh + res (低位乘法再累加到 x1)
        "mul x0, x4, %[prHigh1]\n\t"
        "add x1, x1, x0\n\t"

        // 计算 x1 * p 存回 x1
        "mul x1, x1, %[p1]\n\t"

        // 计算 a[1].low - x1, 存回 x1
        "sub x1, x3, x1\n\t"

        // 计算 x1 和 2 * p 的比较，如果直接减发生溢出就 set flag
        "subs x0, x1, %[twop1]\n\t"
        "csel x1, x0, x1, hs\n\t"

        // 计算 x1 和 p 的比较，如果直接减发生溢出就 set flag
        "subs x0, x1, %[p1]\n\t"
        "csel %[res1], x0, x1, hs\n\t"

        : [ res1 ] "=r"(res1)
        : [ a1 ] "r"(vgetq_lane_u64(a, 1)), [ b1 ] "r"(vgetq_lane_u64(b, 1)), [ p1 ] "r"(vgetq_lane_u64(p, 1)),
        [ twop1 ] "r"(vgetq_lane_u64(twop, 1)), [ prHigh1 ] "r"(vgetq_lane_u64(prHigh, 1)),
        [ prLow1 ] "r"(vgetq_lane_u64(prLow, 1))
        : "x0", "x1", "x3", "x4", "cc");

    // 将结果存储回 r 向量
    *r = vsetq_lane_u64(res0, *r, 0);
    *r = vsetq_lane_u64(res1, *r, 1);
}

/*
 * @brief ：对两个64位整数进行加法运算，并对结果进行模运算
 *
 * @param r：存储结果的指针
 * @param a：第一个操作数
 * @param b：第二个操作数
 * @param p：模数
 */
inline void AddMod_NEON(uint64x2_t *r, uint64x2_t a, uint64x2_t b, uint64x2_t p)
{
    uint64x2_t sum = vaddq_u64(a, b);
    uint64x2_t cmp = vcgtq_u64(sum, p);
    *r = vsubq_u64(sum, vandq_u64(cmp, p));
}

/*
 * @brief ：对两个64位整数进行乘法运算，应用Shoup优化技术并对结果进行模运算
 *
 * @param r：存储结果的指针
 * @param a：第一个操作数
 * @param b：第二个操作数
 * @param bShoup：Shoup优化参数
 * @param q：模数
 */
inline void ShoupMul_NEON(uint64x2_t *r, uint64x2_t a, uint64x2_t b, uint64x2_t bShoup, uint64x2_t q)
{
    uint64_t res0 = 0, res1 = 0;

    // 处理第一个元素
    asm volatile("mul    x3, %[a0], %[bShoup0] \n\t" // x3 = a[0] * bShoup[0] (low 64 bits)
        "umulh  x4, %[a0], %[bShoup0] \n\t"          // x4 = a[0] * bShoup[0] (high 64 bits)
        "mul    x1, x4, %[q0]        \n\t"           // x1 = (a[0] * bShoup[0] high) * q[0]
        "mul    x0, %[a0], %[b0]     \n\t"           // x0 = a[0] * b[0] (low 64 bits)
        "sub    x1, x0, x1           \n\t"           // x1 = (a[0] * b[0]) - ((a[0] * bShoup[0] high) * q[0])
        "subs   x0, x1, %[q0]        \n\t"           // x0 = x1 - q[0]
        "csel   %[res0], x1, x0, lo  \n\t"           // res0 = (x1 < q[0]) ? x1 : x0
        : [ res0 ] "=r"(res0)
        : [ a0 ] "r"(vgetq_lane_u64(a, 0)), [ b0 ] "r"(vgetq_lane_u64(b, 0)),
        [ bShoup0 ] "r"(vgetq_lane_u64(bShoup, 0)), [ q0 ] "r"(vgetq_lane_u64(q, 0))
        : "x0", "x1", "x3", "x4", "cc");

    // 处理第二个元素
    asm volatile("mul    x3, %[a1], %[bShoup1] \n\t" // x3 = a[1] * bShoup[1] (low 64 bits)
        "umulh  x4, %[a1], %[bShoup1] \n\t"          // x4 = a[1] * bShoup[1] (high 64 bits)
        "mul    x1, x4, %[q1]        \n\t"           // x1 = (a[1] * bShoup[1] high) * q[1]
        "mul    x0, %[a1], %[b1]     \n\t"           // x0 = a[1] * b[1] (low 64 bits)
        "sub    x1, x0, x1           \n\t"           // x1 = (a[1] * b[1]) - ((a[1] * bShoup[1] high) * q[1])
        "subs   x0, x1, %[q1]        \n\t"           // x0 = x1 - q[1]
        "csel   %[res1], x1, x0, lo  \n\t"           // res1 = (x1 < q[1]) ? x1 : x0
        : [ res1 ] "=r"(res1)
        : [ a1 ] "r"(vgetq_lane_u64(a, 1)), [ b1 ] "r"(vgetq_lane_u64(b, 1)),
        [ bShoup1 ] "r"(vgetq_lane_u64(bShoup, 1)), [ q1 ] "r"(vgetq_lane_u64(q, 1))
        : "x0", "x1", "x3", "x4", "cc");

    // 将结果存储回 r 向量
    *r = vsetq_lane_u64(res0, *r, 0);
    *r = vsetq_lane_u64(res1, *r, 1);
}

#endif
